package com.qiangesoft.langchain.controller;

import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.SystemMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiStreamingChatModel;
import jakarta.servlet.http.HttpServletResponse;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.List;

@RequestMapping("/ai")
@RestController
public class AiController {

    @Autowired
    private ChatLanguageModel chatLanguageModel;
    @Autowired
    private OpenAiStreamingChatModel openAiStreamingChatModel;

    @GetMapping("/java")
    public String java() {
        OpenAiChatModel chatModel = OpenAiChatModel.builder()
                .apiKey("sk-8d94194115df4559bd318f2e404d4d7e")
                .baseUrl("https://api.deepseek.com/v1")
//                .baseUrl("http://langchain4j.dev/demo/openai/v1")
                .modelName("deepseek-chat")
                .build();
        return chatModel.chat("你好，你是谁？");
    }

    @GetMapping("/boot")
    public String boot() {
        return chatLanguageModel.chat("你好，你是谁？");
    }

    @GetMapping("/dentist")
    public String dentist(@RequestParam(value = "message") String message) {
        List<ChatMessage> messages = new ArrayList<>();
        ChatMessage systemMessage = new SystemMessage("你现在是一个口腔医生，针对患者提出的口腔相关的问题做出回答，超出专业范畴不做回答");
        messages.add(systemMessage);

        ChatMessage userMessage = new UserMessage(message);
        messages.add(userMessage);

        ChatResponse chat = chatLanguageModel.chat(messages);
        return chat.aiMessage().text();
    }

    @GetMapping(value = "/flux", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> flux(@RequestParam(value = "message") String message, HttpServletResponse response) {
        response.setCharacterEncoding("UTF-8");
        return Flux.create(stringFluxSink -> {
            openAiStreamingChatModel.chat(message, new StreamingChatResponseHandler() {
                @Override
                public void onPartialResponse(String s) {
                    stringFluxSink.next(s);
                }

                @Override
                public void onCompleteResponse(ChatResponse chatResponse) {
                    stringFluxSink.complete();
                }

                @Override
                public void onError(Throwable throwable) {
                    stringFluxSink.error(throwable);
                }
            });
        });
    }

}
